# Copyright (C) 2007-2008, Open Society Institute
#                Thomas Moroz: tmoroz@sorosny.org
#
# This program is free software; you can redistribute it and/or modify it
# under the terms of the GNU General Public License Version 2 as published
# by the Free Software Foundation.  You may not use, modify or distribute
# this program under any other version of the GNU General Public License.
#
# This program is distributed in the hope that it will be useful, but
# WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
# General Public License for more details.
# 
# You should have received a copy of the GNU General Public License along
# with this program; if not, write to the Free Software Foundation, Inc.,
# 59 Temple Place - Suite 330, Boston, MA 02111-1307, USA.

import os
import shutil
import tempfile
import cPickle
import logging

from threading import local, Lock
from enfold.xapian.store import Store

logger = logging.getLogger('enfold.indexing.queue')

def get_local(ob, name, factory):
    got = getattr(ob, name, None)
    if got is None:
        setattr(ob, name, factory())
        return getattr(ob, name)
    return got

def make_data_dir(base_dir):
    def make(base_dir=base_dir):
        data_dir = tempfile.mkdtemp(dir=base_dir)
        logger.debug('Will use %r for temporary '
                     'indexing data storage', data_dir)
        return data_dir
    return make

DEL, ADD = range(1, 3)

class QueueItem(object):
    
    def __init__(self, data_dir, data):
        (fd, name) = tempfile.mkstemp(dir=data_dir)
        self.name = name
        f = os.fdopen(fd, 'wb')
        self.dump(data, f)
        
    def update(self, data):
        orig = self.load()
        orig.update(data)

        f = open(self.name, 'wb')
        self.dump(orig, f)

    def dump(self, data, f):
        p = cPickle.Pickler(f, protocol=-1)
        try:
            p.dump(data)
        finally:
            del p
            f.flush()
            f.close()

    def load(self):
        f = open(self.name, 'rb')
        p = cPickle.Unpickler(f)
        try:
            orig = p.load()
        finally:
            del p
            f.close()
        return orig

    def remove(self):
        os.remove(self.name)

class QueueStore(Store):
    """A thread-safe queue store
    """

    def __init__(self, *args, **kw):
        self.base_dir = kw.get('base_dir')
        self.add_hook = kw.get('add_hook')
        self.remove_hook = kw.get('remove_hook')
        self.search_hook = kw.get('search_hook')
        for name in ('add_hook', 'remove_hook', 'search_hook', 'base_dir'):
            try:
                del kw[name]
            except KeyError:
                pass
        Store.__init__(self, *args, **kw)
        self.lock = Lock()
        self._reset()

    def setHooks(self, add_hook=None, remove_hook=None, search_hook=None):
        self.add_hook = add_hook
        self.remove_hook = remove_hook
        self.search_hook = search_hook

    def add(self, data):
        data = dict(data) # Copy.
        uid = data.get('uid', None)

        if not uid:
            msg = 'You must pass in the uid of the object to be indexed.'
            raise ValueError, msg

        if self.add_hook is not None:
            self.add_hook()

        info = get_local(self.queue_local, 'info', dict)
        queue = get_local(self.queue_local, 'queue', list)
        data_dir = get_local(self.queue_local, 'data_dir', 
                             make_data_dir(self.base_dir))
        
        if uid in info:
            info[uid].update(data)
        else:
            info[uid] = QueueItem(data_dir, data)
        queue.append((uid, ADD))
        self.queue_local.isEmpty = False

    def remove(self, uid):
        if not uid:
            msg = 'You must pass in the uid of the object to be indexed.'
            raise ValueError, msg

        if self.remove_hook is not None:
            self.remove_hook()

        queue = get_local(self.queue_local, 'queue', list)
        queue.append((uid, DEL))
        self.queue_local.isEmpty = False

    def isEmpty(self):
        return get_local(self.queue_local, 'isEmpty', lambda: True)

    def _reset(self):
        self.queue_local = local()
        
    def reset(self, *args, **kw):
        self._reset()
        Store.reset(self, *args, **kw)

    def prepare(self):
        if self.isEmpty():
            return

        queue = get_local(self.queue_local, 'queue', list)
        info = get_local(self.queue_local, 'info', dict)
        ops = get_local(self.queue_local, 'ops', dict)
        pieces = get_local(self.queue_local, 'pieces', dict)
        data_dir = get_local(self.queue_local, 'data_dir', 
                             make_data_dir(self.base_dir))

        # Weed out operations that cancel one another.
        #
        # Rules:
        #
        # 1. A remove after an add for the same uid means we
        #    don't even try to add it.
        # 2. For an add after a remove, we skip the remove since
        #    Xapian will 'replace' the document under the hood.
        # 3. For one or more adds, ignore all but the last one.
        #
        # With some logic, we can reduce those rules to just one:
        #
        #  "Whatever the last operation is, just use that
        #   and ignore everything else."

        for idx, (uid, op) in enumerate(queue):
            ops[(uid, op)] = (idx, uid, op)
            
        # Pre-process data into DocumentPieces. This is done before
        # calling addDocument to keep the connection open for as
        # little as possible.
        for idx, uid, op in sorted(ops.itervalues()):
            if op in (ADD,):
                item = info[uid]
                data = self.prepare_data(item.load())
                info[uid] = QueueItem(data_dir, data)
                item.remove()
                pieces[uid] = QueueItem(data_dir, self.process(data))

    def processQueue(self, clear=True):
        info = get_local(self.queue_local, 'info', dict)
        ops = get_local(self.queue_local, 'ops', dict)
        pieces = get_local(self.queue_local, 'pieces', dict)

        processed = False

        if self.isEmpty():
            return processed

        # Now process the queue.
        for idx, uid, op in sorted(ops.itervalues()):
            processed = True
            if op in (DEL,):
                Store.remove(self, uid)
            elif op in (ADD,):
                i = info.get(uid)
                if i is not None:
                    i = i.load()
                p = pieces.get(uid)
                if p is not None:
                    p = p.load()
                Store.add(self, i, p, prepare=False)
            else:
                raise ValueError, 'processQueue() encountered unknown operator.'

        if clear:
            self.clearQueue()
        return processed

    def clearQueue(self):
        if self.isEmpty():
            return

        queue = get_local(self.queue_local, 'queue', list)
        del queue[:]
        
        info = get_local(self.queue_local, 'info', dict)
        info.clear()

        ops = get_local(self.queue_local, 'ops', dict)
        ops.clear()

        pieces = get_local(self.queue_local, 'pieces', dict)
        pieces.clear()

        data_dir = get_local(self.queue_local, 'data_dir', 
                             make_data_dir(self.base_dir))

        shutil.rmtree(data_dir)

        self.queue_local.data_dir = None
        self.queue_local.isEmpty = True

    def getQueueState(self):
        info = get_local(self.queue_local, 'info', dict)
        queue = get_local(self.queue_local, 'queue', list)
        return info.copy(), queue[:]

    def setQueueState(self, state):
        info, queue = state
        self.queue_local.info = info
        self.queue_local.queue = queue
        if queue:
            self.queue_local.isEmpty = False
        else:
            self.queue_local.isEmpty = True

    def search(self, *args, **kw):
        if self.search_hook is not None:
            self.search_hook()
        return Store.search(self, *args, **kw)
